import sys
import numpy as np
from PySide6.QtWidgets import (
    QApplication, QGraphicsView, QGraphicsScene, QGraphicsEllipseItem, QGraphicsSimpleTextItem, QVBoxLayout, QWidget
)
from PySide6.QtGui import QPen, QColor, QBrush, QPainter
from PySide6.QtCore import Qt, QPointF

# 导入必要的 PySide6 模块 (用于GUI) 和 numpy (用于数据处理)

# 模拟数据
np.random.seed(0)
num_points = 30
well_id = 'A1'
depths = np.linspace(1000, 2000, num_points)
rock_desc = ['砂岩', '泥岩', '灰岩']
analyst = '张三'
date = '2024-06-01'
types = ['a', 'b', 'c']
data = []
# 定义模拟所需的基本参数：数据点数量、井号、深度范围、岩性、分析人、日期、类型
for i in range(num_points):
    row = {
        '井号': well_id,
        '深度': depths[i],
        '岩性简述': np.random.choice(rock_desc),
        '分析人': analyst,
        '分析日期': date,
        '类型': np.random.choice(types),
        'Mg': np.random.uniform(0, 10),
        'Al': np.random.uniform(0, 10),
        'Si': np.random.uniform(0, 10),
        'P': np.random.uniform(0, 10),
        'S': np.random.uniform(0, 10),
        'Cl': np.random.uniform(0, 10),
        'K': np.random.uniform(0, 10),
        'Ca': np.random.uniform(0, 10),
        'Ti': np.random.uniform(0, 10),
        'V': np.random.uniform(0, 10),
        'Mn': np.random.uniform(0, 10),
        'Fe': np.random.uniform(0, 10),
    }
    data.append(row)

# 循环生成指定数量 (num_points) 的模拟数据点，每个点包含各种字段信息和随机生成的元素含量

def get_line_data(data, element):
    return [row[element] for row in data]

def get_depths(data):
    return [row['深度'] for row in data]

# 定义两个辅助函数，用于从数据列表中提取特定元素的含量列表和深度列表

class DraggableView(QGraphicsView):
    def __init__(self, scene):
        super().__init__(scene)
        self.setRenderHint(QPainter.Antialiasing)
        self.setDragMode(QGraphicsView.ScrollHandDrag)
        self.setMouseTracking(True)
        self.tooltip = QGraphicsSimpleTextItem()
        self.tooltip.setZValue(10)
        self.scene().addItem(self.tooltip)
        self.tooltip.hide()

    def mouseMoveEvent(self, event):
        pos = self.mapToScene(event.pos())
        items = self.scene().items(QPointF(pos.x(), pos.y()))
        for item in items:
            if isinstance(item, DataPointItem):
                self.tooltip.setText(item.data_text)
                self.tooltip.setPos(pos.x() + 10, pos.y() - 10)
                self.tooltip.show()
                break
        else:
            self.tooltip.hide()
        super().mouseMoveEvent(event)

# 自定义的 QGraphicsView 类，实现鼠标拖动平移场景和鼠标跟踪显示数据点的 tooltip

class DataPointItem(QGraphicsEllipseItem):
    def __init__(self, x, y, data_text, color):
        super().__init__(x-4, y-4, 8, 8)
        self.setBrush(QBrush(color))
        self.setPen(QPen(Qt.black))
        self.setZValue(5)
        self.data_text = data_text
        self.setAcceptHoverEvents(True)

    def hoverEnterEvent(self, event):
        self.setBrush(QBrush(QColor(255, 0, 0)))
        super().hoverEnterEvent(event)

    def hoverLeaveEvent(self, event):
        self.setBrush(QBrush(QColor(0, 150, 255)))
        super().hoverLeaveEvent(event)

# 自定义的 QGraphicsEllipseItem 类，表示图表上的一个数据点，实现鼠标悬停高亮和存储 tooltip 文本

class PlotWidget(QWidget):
    def __init__(self, data):
        super().__init__()
        layout = QVBoxLayout(self)
        self.scene = QGraphicsScene(self)
        self.view = DraggableView(self.scene)
        layout.addWidget(self.view)
        self.setLayout(layout)
        self.draw_plot(data)

    def draw_plot(self, data):
        # 选择三个元素作为a/b/c
        elements = ['Mg', 'Al', 'Si']
        colors = [QColor(255, 0, 0), QColor(0, 150, 255), QColor(0, 200, 0)]
        depths = get_depths(data)
        min_depth, max_depth = min(depths), max(depths)
        margin = 60
        width, height = 900, 600
        plot_height = height - 2 * margin
        plot_width = width - 2 * margin
        # 横坐标为深度
        min_x, max_x = min_depth, max_depth
        # 纵坐标为元素含量
        min_y, max_y = 0, 12
        # 画坐标轴
        self.scene.addLine(margin, margin, margin, height - margin, QPen(Qt.black))
        self.scene.addLine(margin, height - margin, width - margin, height - margin, QPen(Qt.black))
        # 画纵坐标刻度（元素含量）
        for i in range(7):
            y = height - margin - i * (plot_height / 6)
            val = min_y + (max_y - min_y) * i / 6
            self.scene.addLine(margin-5, y, margin, y, QPen(Qt.black))
            text = QGraphicsSimpleTextItem(f"{val:.1f}")
            text.setPos(10, y-10)
            self.scene.addItem(text)
        # 画横坐标刻度（深度）
        for i in range(6):
            x = margin + i * (plot_width / 5)
            depth_val = min_x + (max_x - min_x) * i / 5
            self.scene.addLine(x, height - margin, x, height - margin + 5, QPen(Qt.black))
            text = QGraphicsSimpleTextItem(f"{depth_val:.0f}")
            text.setPos(x-20, height - margin + 10)
            self.scene.addItem(text)
        # 画折线和点
        for idx, element in enumerate(elements):
            values = get_line_data(data, element)
            points = []
            for i, (depth, val) in enumerate(zip(depths, values)):
                x = margin + (depth - min_x) / (max_x - min_x) * plot_width
                y = height - margin - (val - min_y) / (max_y - min_y) * plot_height
                points.append((x, y))
                data_text = f"类型: {['a','b','c'][idx]}\n深度: {depth:.1f}\n{element}: {val:.2f}"
                pt = DataPointItem(x, y, data_text, colors[idx])
                self.scene.addItem(pt)
            # 画折线
            for i in range(len(points)-1):
                self.scene.addLine(points[i][0], points[i][1], points[i+1][0], points[i+1][1], QPen(colors[idx], 2))
        # 图例
        for idx, element in enumerate(elements):
            legend = QGraphicsSimpleTextItem(f"{['a','b','c'][idx]}: {element}")
            legend.setPos(width - margin - 80, margin + idx*20)
            legend.setBrush(QBrush(colors[idx]))
            self.scene.addItem(legend)
        self.scene.setSceneRect(0, 0, width, height)

# 主窗口部件，包含 QGraphicsScene 和 DraggableView，并负责绘制图表元素 (坐标轴、刻度、折线、点、图例)

def main():
    app = QApplication(sys.argv)
    w = PlotWidget(data)
    w.setWindowTitle('井数据可视化-地质导向交互折线图')
    w.resize(1000, 700)
    w.show()
    sys.exit(app.exec())

# 主函数，初始化 QApplication，创建 PlotWidget 窗口并显示，进入应用事件循环

if __name__ == '__main__':
    main()

# 程序入口点，判断当前脚本是否作为主程序运行，如果是则调用 main() 函数启动应用
